#include <libipq.h>
#include <linux/netfilter.h>
#include <stdio.h>
#include <string.h>
#include <netinet/tcp.h>
#include <netinet/ip.h>

#include <stdlib.h>

#include "icmp_tt.c"

#define BUFSIZE 2048
#define UDP	17
#define TCP	6
#define ICMP	1

#define MAX_PING_SESSIONS 1000

#define ICMP_ECHO_REPLY 000
#define ICMP_ECHO_REQ 010

struct icmp_hdr {
	char type;
	char code;
	short checksum;
	short identifier;
	short seq_num;
} __attribute__((__packed__));

struct udp_hdr {
	short src_port;
	short dst_port;
	short length;
	short checksum;
} __attribute__((__packed__));


static void die(struct ipq_handle *h) {
	ipq_perror("passer");
	ipq_destroy_handle(h);
	exit(1);
}

void change_ip_addr(struct iphdr* header, struct in_addr* new_addr) {
	struct in_addr* src_addr;

	src_addr = (struct in_addr*) &header->saddr;
	*src_addr = *new_addr;
}

int handle_icmp(ipq_packet_msg_t* m, struct iphdr* ip, void* icmp_tt) {

	int status;
	char identifier[80] = {'\0'};
	char seq_num[80] = {'\0'};
	char* hash_key;
	struct in_addr* orig_src_addr;
	struct icmp_hdr *icmp = (struct icmp_hdr*)(m->payload + (4 * ip->ihl));

	//Create the first part of hashkey
	memset(identifier, '\0', 80);
	memset(seq_num, '\0', 80);
	hash_key = calloc(1, 160);
	sprintf(identifier, " %x ", ntohs(icmp->identifier));
	hash_key = strcat(hash_key, identifier);
	sprintf(seq_num, " %hd ", ntohs(icmp->seq_num));
	hash_key = strcat(hash_key, seq_num);

	if (icmp->type == ICMP_ECHO_REQ) {
		printf("Icmp echo\n");
		//Second part of hashkey
		hash_key = strcat(hash_key, (char* ) inet_ntoa(*(struct in_addr*) &ip->daddr));
		printf("Hashkey: %s\n", hash_key);

		//Add to icmp translation table
		insert(hash_key, malloc(sizeof(struct in_addr)), icmp_tt);
		orig_src_addr = (struct in_addr*) search(hash_key, icmp_tt);
		memcpy(orig_src_addr, &ip->saddr, sizeof(struct in_addr));

		//Change source ip to NAT


		status = NF_ACCEPT;

	} else if (icmp->type == ICMP_ECHO_REPLY) {
		printf("Icmp reply\n");
		//Second part of hashkey
		hash_key = strcat(hash_key, (char* ) inet_ntoa(*(struct in_addr*) &ip->saddr));
		printf("Hashkey: %s\n", hash_key);

		//Look from translation table
		orig_src_addr = (struct in_addr*) search(hash_key, icmp_tt);
		if (orig_src_addr != NULL) {
			printf("Session found %s\n", (char* )inet_ntoa(*orig_src_addr));
			//Change source ip to host

			//Delete session
			if (delete_item(hash_key, icmp_tt) == 1) {
				free(orig_src_addr);
				free(hash_key);
				printf("Session deleted\n");
			}
		} else {
			printf("No corresponding ping session found...\n");
		}
		status = NF_ACCEPT;
	} else {
		printf("Other icmp message\n");
		status = NF_DROP;
	}

	return status;
}

int handle_udp(ipq_packet_msg_t* m, struct iphdr* ip, void* udp_tt){
	
	return NF_ACCEPT;


	char* hash_key;
	char port[80] = {'\0'};
	short* original_port;
	int status = NF_DROP;
	int new_port;
	
	struct udp_hdr *udp = (struct udp_hdr*)(m->payload + (4 * ip->ihl));	
	
	//Create the first part of hashkey
	hash_key = calloc(1, 160);
	
	if(1/*IS_OUTBOUND*/){
		//Get and set new port for packet	
		new_port = 55000;
		
		//Create hashkey
		sprintf(port, " %x ", new_port);
		hash_key = strcat(hash_key, port);
		hash_key = strcat(hash_key, (char* )inet_ntoa(*(struct in_addr*) &ip->daddr));

		//Add port&ip to udp_tt
		original_port = calloc(1, sizeof(short));
		*original_port = ntohs(udp->src_port);
		insert(hash_key, (char* )original_port, udp_tt);
		//Forward
		status = NF_ACCEPT;
	}else if(0/*IS_INBOUND*/){
		//Create hashkey
		sprintf(port, " %x ", ntohs(udp->src_port));
		hash_key = strcat(hash_key, port);
		hash_key = strcat(hash_key, (char* )inet_ntoa(*(struct in_addr*) &ip->daddr));
		
		//Check port&ip from udp_tt
		original_port = (short* ) search(hash_key, udp_tt);
		if(original_port == NULL){
			//No connection found => DROP
			status = NF_DROP;
		}else{
			//Do port and address translation			
			udp->dst_port = *original_port;
			//*ip->saddr = portaddrstruct->orig_addr;
			status = NF_ACCEPT;
		}

	}else{
		//???
	}
	
	
	return status;
}

int handle_tcp(){
	
	return NF_ACCEPT;
}

int main(int argc, char **argv) {
	int status;
	unsigned char buf[BUFSIZE];
	struct ipq_handle *h;
	int packet_status;

	void* icmp_tt = create_list();
	void* udp_tt = create_list();

	h = ipq_create_handle(0, PF_INET);
	if (!h)
		die(h);

	status = ipq_set_mode(h, IPQ_COPY_PACKET, BUFSIZE);
	if (status < 0)
		die(h);

	do {
		status = ipq_read(h, buf, BUFSIZE, 0);
		if (status < 0)
			die(h);

		switch (ipq_message_type(buf)) {
		case NLMSG_ERROR:
			fprintf(stderr, "Received error message %d\n", ipq_get_msgerr(buf));
			break;

		case IPQM_PACKET: {
				ipq_packet_msg_t *m = ipq_get_packet(buf);

				struct iphdr *ip = (struct iphdr*) m->payload;

				switch (ip->protocol) {

				case UDP:
					printf("UDP PACKET\n");
					packet_status = handle_udp(m, ip, udp_tt);
					break;

				case TCP:
					printf("TCP PACKET\n");
					packet_status = handle_tcp();
					break;

				case ICMP:
					printf("ICMP PACKET\n");
					packet_status = handle_icmp(m, ip, icmp_tt);
					break;

				default:
					break;
				}
				
				status = ipq_set_verdict(h, m->packet_id, packet_status, 0, NULL);


				if (status < 0)
					die(h);
				break;
			}

		default:
			fprintf(stderr, "Unknown message type!\n");
			break;
		}
	} while (1);

	ipq_destroy_handle(h);
	return 0;
}